import os
import re
import cv2
from pathlib import Path
from typing import Optional, Tuple, Union

PathLike = Union[str, Path, os.PathLike[str]]

def images_to_video(
    dir_path: PathLike,
    pattern: str = "*.png",
    output_name: Optional[str] = None,
    fps: int = 30,
    codec: str = "mp4v",
    recursive: bool = False,
    resize_to: Optional[Tuple[int, int]] = None,
    strict_same_size: bool = False,
) -> str:
    """
    将目录中的一批图片拼成视频。返回生成的视频绝对路径。
    """
    p = Path(dir_path)  # 不改变参数本身的类型，避免类型检查告警
    if not p.is_dir():
        raise NotADirectoryError(f"目录不存在: {p}")

    # 用 Path.glob / rglob 取文件列表，而不是拼字符串
    it = p.rglob(pattern) if recursive else p.glob(pattern)
    files = [str(x) for x in it]

    if not files:
        raise RuntimeError(f"未找到匹配的图片: {p} / {pattern}")

    # 自然排序：1,2,10 按数字顺序
    def nat_key(s: str):
        return [int(t) if t.isdigit() else t.lower() for t in re.split(r"(\d+)", Path(s).name)]
    files.sort(key=nat_key)

    # 输出路径
    if output_name is None:
        output_name = p.name + ".mp4"
    out_path = str(p / output_name)

    # 首帧与尺寸
    frame0 = cv2.imread(files[0], cv2.IMREAD_UNCHANGED)
    if frame0 is None:
        raise RuntimeError(f"无法读取图片: {files[0]}")
    if frame0.ndim == 3 and frame0.shape[2] == 4:
        frame0 = cv2.cvtColor(frame0, cv2.COLOR_BGRA2BGR)

    if resize_to is not None:
        w, h = resize_to
        frame0 = cv2.resize(frame0, (w, h), interpolation=cv2.INTER_AREA)
    else:
        h, w = frame0.shape[:2]

    writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*codec), fps, (w, h))
    # 兼容老一点的 OpenCV：用构造方式
    writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*codec), fps, (w, h))
    if not writer.isOpened():
        raise RuntimeError("VideoWriter 打开失败，请检查编码器/路径是否可写")

    written = 0
    for f in files:
        img = cv2.imread(f, cv2.IMREAD_UNCHANGED)
        if img is None:
            print(f"[跳过不可读] {f}")
            continue
        if img.ndim == 3 and img.shape[2] == 4:
            img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
        if img.shape[0] != h or img.shape[1] != w:
            if strict_same_size:
                raise ValueError(f"尺寸不一致: {f} 是 {img.shape[1]}x{img.shape[0]}，目标是 {w}x{h}")
            img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
        writer.write(img)
        written += 1

    writer.release()
    if written == 0:
        raise RuntimeError("没有成功写入任何帧")
    print(f"完成 -> {out_path}  | 帧数: {written} | fps: {fps} | 尺寸: {w}x{h}")
    return out_path


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--dir", type=str, help="图片目录")
    parser.add_argument("--pattern", type=str, default="*.png", help="图片文件名模式")
    parser.add_argument("--output-name", type=str, default=None, help="输出视频文件名，默认是目录名.mp4")
    parser.add_argument("--fps", type=int, default=30, help="视频帧率")
    parser.add_argument("--codec", type=str, default="mp4v", help="视频编码器，如 mp4v, X264, avc1, mjpg 等")
    parser.add_argument("--recursive", action="store_true", help="递归子目录查找图片")
    parser.add_argument("--resize-to", type=int, nargs=2, metavar=("W", "H"), default=None,
                        help="调整所有图片到指定尺寸 (宽 高)，默认不调整")
    parser.add_argument("--strict-same-size", action="store_true",
                        help="如果有图片尺寸与首张不符则报错，默认是自动调整尺寸")

    args = parser.parse_args()
    images_to_video(
        dir_path="D:\\workspace\\play_with_pybullet\\pic",  # args.dir,
        pattern=args.pattern,
        output_name=args.output_name,
        fps=args.fps,
        codec=args.codec,
        recursive=args.recursive,
        resize_to=tuple(args.resize_to) if args.resize_to else None,
        strict_same_size=args.strict_same_size,
    )